# -*- coding: utf-8 -*-
import networkx as nx
from node2vec import Node2Vec

# Create a graph
graph = nx.fast_gnp_random_graph(n=3, p=0.5)

# Precompute probabilities and generate walks
node2vec = Node2Vec(graph, dimensions=64, walk_length=30, num_walks=200, workers=1)

# Embed
model = node2vec.fit(window=10, min_count=1, batch_words=4)  # Any keywords acceptable by gensim.Word2Vec can be passed, `diemnsions` and `workers` are automatically passed (from the Node2Vec constructor)

# Look for most similar nodes
print model.wv.most_similar('1')  # Output node names are always strings

# Save embeddings for later use
model.wv.save_word2vec_format("aa")

# Save model for later use
model.save("bb")
